#include "accel.h"

namespace csupbr {

    struct BVHPrimitiveInfo {
		size_t primitiveNumber = 0;
		Point3f centroid;
		BBox3f bbox;
		BVHPrimitiveInfo() = default;
		BVHPrimitiveInfo(size_t i, const BBox3f& bbox)
			:primitiveNumber(i), bbox(bbox) {
			centroid = .5f * bbox.pMin + .5f * bbox.pMax;
		}
	};

	struct BVHBuildNode {
		BBox3f bbox;
		union {
			BVHBuildNode* children[2];
			struct {
				BVHBuildNode* left, * right;
			};
		};
		int splitAxis = 0, firstPrimOffset = 0, nPrimitives = 0;
		void initLeaf(int firstOffset, int nPrim, const BBox3f& bbox) {
			firstPrimOffset = firstOffset;
			nPrimitives = nPrim;
			this->bbox = bbox;
			left = right = nullptr;
		}
		void initInterior(int axis, BVHBuildNode* c0, BVHBuildNode* c1) {
			left = c0;
			right = c1;
			bbox = BBox3f::Union(c0->bbox, c1->bbox);
			splitAxis = axis;
			nPrimitives = 0;
		}
	};

    void BVHTree::build() {
        if (primitives.size() == 0) { return; }
		// Build BVH from primitives
		std::vector<BVHPrimitiveInfo> primitiveInfo;
		primitiveInfo.reserve(primitives.size());
		for (int i = 0; i < primitives.size(); i++) {
			primitiveInfo.emplace_back(i, dynamic_cast<GeoPrimitive*>(primitives[i].get())->getBBox());
		}
		int totalNodes = 0;
		std::vector<std::shared_ptr<Primitive>>orderedPrims;
		orderedPrims.reserve(primitives.size());
		BVHBuildNode* root;
		root = recursiveBuild(primitiveInfo, 0, primitiveInfo.size(), totalNodes, orderedPrims);
		this->primitives.swap(orderedPrims);
		primitiveInfo.resize(0);
		lbnodes.resize(totalNodes);
		int offset = 0;
		flattenTree(root, offset);
		assert(offset == totalNodes);
    }

    BVHBuildNode* BVHTree::recursiveBuild(std::vector<BVHPrimitiveInfo>& primitiveInfo,
		int start, int end, int& totalNodes, std::vector<std::shared_ptr<Primitive>>& orderedPrims) {
		BVHBuildNode* root = new BVHBuildNode();
		++totalNodes;
		BBox3f bbox;
		for (int i = start; i < end; i++)
			bbox = BBox3f::Union(bbox, primitiveInfo[i].bbox);
		int nPrimitives = end - start;
		if (nPrimitives == 1) {
			int fisrtPrimOffset = orderedPrims.size();
			for (int i = start; i < end; ++i) {
				int primNum = primitiveInfo[i].primitiveNumber;
				orderedPrims.push_back(primitives[primNum]);
			}
			root->initLeaf(fisrtPrimOffset, nPrimitives, bbox);
			return root;
		}
		else {
			// Compute bbox of primitive centriod, choose split dimension 
			BBox3f centroidBox;
			for (int i = start; i < end; ++i)
				centroidBox = BBox3f::Union(centroidBox, primitiveInfo[i].centroid);
			int dim = centroidBox.maxExtent();
			// Partition primitives into two sets and build children
			int mid = (start + end) / 2;
			if (centroidBox.pMax[dim] == centroidBox.pMin[dim]) { // Centroid bbox degenerate
				// Create leaf node
				int fisrtPrimOffset = orderedPrims.size();
				for (int i = start; i < end; ++i) {
					int primNum = primitiveInfo[i].primitiveNumber;
					orderedPrims.push_back(primitives[primNum]);
				}
				root->initLeaf(fisrtPrimOffset, nPrimitives, bbox);
				return root;
			}
			else {
				// Partition primitives based on splitmethod
				mid = (start + end) / 2;
                std::nth_element(&primitiveInfo[start], &primitiveInfo[mid],
                    &primitiveInfo[(size_t)end - 1] + 1,
                    [dim](const BVHPrimitiveInfo& a, const BVHPrimitiveInfo& b) {
                        return a.centroid[dim] < b.centroid[dim];
                    });
				root->initInterior(
					dim,
					recursiveBuild(primitiveInfo, start, mid, totalNodes, orderedPrims),
					recursiveBuild(primitiveInfo, mid, end, totalNodes, orderedPrims)
				);
			}
		}
		return root;
	}

	int BVHTree::flattenTree(BVHBuildNode* node, int& offset) {
		LinearBVHNode& linearNode = lbnodes[offset];
		linearNode.bbox = node->bbox;
		int myOffset = offset++;
		if (node->nPrimitives > 0) {
			linearNode.primitiveOffset = node->firstPrimOffset;
			linearNode.nPrimitives = node->nPrimitives;
		}
		else {
			// Create interior flattened BVH node
			linearNode.axis = node->splitAxis;
			linearNode.nPrimitives = 0;
			flattenTree(node->children[0], offset);
			linearNode.secondChildOffset =
				flattenTree(node->children[1], offset);
		}
		return myOffset;
	}

    bool BVHTree::intersect(const Ray &ray, Intersection &isct) const {
        if (lbnodes.empty()) return false;
		bool hit = false;
		Vec3f invDir(1 / ray.direct.x, 1 / ray.direct.y, 1 / ray.direct.z);
		int dirIsNeg[3] = { invDir.x < 0, invDir.y < 0, invDir.z < 0 };
		// Follow ray through BVH nodes to find primitive intersections
		int toVisitOffset = 0, currentNodeIndex = 0;
		int nodesToVisit[64];
		while (true) {
			const LinearBVHNode& node = lbnodes[currentNodeIndex];
			// Check ray against BVH node
			float hit0, hit1;
			if (node.bbox.intersectP(ray, hit0, hit1)) {
				if (node.nPrimitives > 0) {
					// Intersect ray with primitives in leaf BVH node
					for (int i = 0; i < node.nPrimitives; ++i)
						if (primitives[node.primitiveOffset + i]->intersect(ray, isct)) {
							hit = true;
						}
					if (toVisitOffset == 0) break;
					currentNodeIndex = nodesToVisit[--toVisitOffset];
				}
				else {
					// Put far BVH node on _nodesToVisit_ stack, advance to near
					// node
					if (dirIsNeg[node.axis]) {
						nodesToVisit[toVisitOffset++] = currentNodeIndex + 1;
						currentNodeIndex = node.secondChildOffset;
					}
					else {
						nodesToVisit[toVisitOffset++] = node.secondChildOffset;
						currentNodeIndex = currentNodeIndex + 1;
					}
				}
			}
			else {
				if (toVisitOffset == 0) break;
				currentNodeIndex = nodesToVisit[--toVisitOffset];
			}
		}
		return hit;
    }

    bool BVHTree::intersectP(const Ray &ray, bool blend_test) const {
        if (lbnodes.empty()) return false;
        Vec3f invDir(1 / ray.direct.x, 1 / ray.direct.y, 1 / ray.direct.z);
        int dirIsNeg[3] = { invDir.x < 0, invDir.y < 0, invDir.z < 0 };
        // Follow ray through BVH nodes to find primitive intersections
        int toVisitOffset = 0, currentNodeIndex = 0;
        int nodesToVisit[64];
        while (true) {
            const LinearBVHNode& node = lbnodes[currentNodeIndex];
            // Check ray against BVH node
            float hit0, hit1;
            if (node.bbox.intersectP(ray, hit0, hit1)) {
                if (node.nPrimitives > 0) {
                    // Intersect ray with primitives in leaf BVH node
                    for (int i = 0; i < node.nPrimitives; ++i)
                        if (primitives[node.primitiveOffset + i]->intersectP(ray, blend_test)) {
                            return true;
                        }
                    if (toVisitOffset == 0) break;
                    currentNodeIndex = nodesToVisit[--toVisitOffset];
                }
                else {
                    // Put far BVH node on _nodesToVisit_ stack, advance to near
                    // node
                    if (dirIsNeg[node.axis]) {
                        nodesToVisit[toVisitOffset++] = currentNodeIndex + 1;
                        currentNodeIndex = node.secondChildOffset;
                    }
                    else {
                        nodesToVisit[toVisitOffset++] = node.secondChildOffset;
                        currentNodeIndex = currentNodeIndex + 1;
                    }
                }
            }
            else {
                if (toVisitOffset == 0) break;
                currentNodeIndex = nodesToVisit[--toVisitOffset];
            }
        }
        return false;
    }
}